import os
from datetime import datetime

from langchain.document_loaders import TextLoader
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS

from es_search import search_and_save_results
from concurrent.futures import ThreadPoolExecutor
executor = ThreadPoolExecutor(max_workers=100)



# 文件路径和模型路径
embedding_path = '/data/model/zpoint_large_embedding_zh'
# 创建线程池执行器

embeddings = HuggingFaceEmbeddings(model_name=embedding_path)

import threading
from contextlib import contextmanager

from langchain.vectorstores.faiss import FAISS


class ThreadSafeFaiss:
    def __init__(self, obj: FAISS):
        self._obj = obj
        self._lock = threading.RLock()

    # 获取线程锁，只有获取到线程锁的 Faiss 对象才能执行操作

    @contextmanager
    def acquire(self):
        try:
            self._lock.acquire()
            yield self._obj
        finally:
            self._lock.release()

    def save_local(self, path: str):
        with self.acquire():
            if not os.path.isdir(path):
                os.makedirs(path)
            self._obj.save_local(path)

    def delete(self):
        ret = []
        with self.acquire():
            ids = list(self._obj.docstore._dict.keys())
            if ids:
                ret = self._obj.delete(ids)
        return ret

def create_docsearch_index(file_path, embedding_path, query):
    try:
        print(file_path)
        if os.path.exists(file_path):
            loader = TextLoader(file_path)
            documents = loader.load()

            text_splitter = RecursiveCharacterTextSplitter(chunk_size=0, chunk_overlap=0, separators=["-----"])
            texts = text_splitter.split_documents(documents)

            docsearch = ThreadSafeFaiss.from_documents(texts, embeddings)
            results = docsearch.similarity_search(query, k=1)
            return "xxxx"
        else:
            print("不存在")


    except Exception as e:
        print("======================")
        print(str(e))
        print()


def prom(content, question, save_path):
    context = content
    prompt_template = (
        "{question}\n"
        "背景知识：\n{context}")
    prompt_text = prompt_template.format(context=context, question=question)
    PROMPT = PromptTemplate(template=prompt_text, input_variables=["context", "question"])
    os.remove(save_path)
    return PROMPT.template


def test(question, startYear, endYear):
    print(question)
    t1 = datetime.now()
    save_path = search_and_save_results(question, startYear, endYear)

    t2 = datetime.now()
    second = (t1 - t2).total_seconds()

    background = create_docsearch_index(save_path, embedding_path, question)
    print(len(background))

    re = "ok"

    return re



threads = []

for x in range(5):
    future=executor.submit(test, "福州经济政策", None, None)
    threads.append(future)

for x in threads:
    print(future.result())